# bge-reranker-large  quietnight/bge-reranker-large
from modelscope import snapshot_download
from BCEmbedding import  RerankerModel
import base_util


class ReRankModel:
    """
    抽象基类，定义了Embedding模型的接口。
    """
    def __init__(self, model_name):
        #
        self.model_download_path=f"{base_util.current_project_dir}\\\models\ReRankModels\\"
        self.model_name=model_name
        self.model=None



    def load_model(self,model_id):
        # model_path=self.download_model(model_id=model_id,cache_dir=self.model_download_path)
        pass

    def download_model(self,model_id,cache_dir):
        model_dir = snapshot_download(model_id, cache_dir=cache_dir)
        return model_dir



class BCEReRankModel(ReRankModel):


    def __init__(self, model_name='maidalun/bce-reranker-base_v1'):
        # super().__init__('BAAI/bge-large-zh-v1.5')
        super().__init__(model_name)
        self.load_model(model_name)

    def load_model(self, model_id='maidalun/bce-reranker-base_v1'):
        model_dir = snapshot_download(model_id, cache_dir=self.model_download_path)
        self.model = RerankerModel(model_name_or_path=model_dir)
        return self.model


    def compute_score(self,query, passages):
        """
         计算分数
        :param sentence_pairs:
        :return:
        """
        sentence_pairs = [[query, passage] for passage in passages]
        scores = self.model.compute_score(sentence_pairs)
        return scores

    def rerank(self,query, passages):
        """
        重排序
        :param query:
        :param passages:
        :return:
        """
        rerank_results = self.model.rerank(query, passages)
        return rerank_results


if __name__ == '__main__':


    pass